import os
import torch
from torchvision import transforms, models
import cv2
from PIL import Image
classes_lj = {"CN":[["CN1","CN2","CN3","CN4","IC1","IC2","IC3","MIS"],"./AI/weights/resNetIC_CN_815.pth"],
              "IC":[["IC1","IC2","IC3","MIS"],"./AI/weights/resNetIC_MIS.pth"]}

classes_indices = { "CN1":["BRIDGE","Lead_LIFT","OK","SHIFT","SOLDER"],
                    "CN3":["BRIDGE","Lead_LIFT","OK","SHIFT","SOLDER"],
                    "CN4":["NG","OK"],
                    "CN4_chajian":["1","2","3","4","5","6"],
                    "IC1":["BRIDGE","OK","SHIFT","SOLDER"],
                    "IC3":["BRIDGE","Lead_LIFT","OK","SHIFT","SOLDER"]}
classes_models = {  "CN1":"./AI/weights/resNetCN1_classification.pth",
                    "CN3":"./AI/weights/resNetCN3_classification.pth",
                    "CN4":"./AI/weights/resNetCN4_classification.pth",
                    "CN4_chajian":"./AI/weights/resNetCN4_chajian.pth",
                    "IC1":"./AI/weights/resNetIC1_classification.pth",
                    "IC3":"./AI/weights/resNetIC3_classification.pth"}

class AI_Classify():
    def __init__(self):
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model = models.resnet34(pretrained=False)

    def date_transform(self,image):

        img_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        data_transform = transforms.Compose(
            [transforms.Resize(256),
             transforms.CenterCrop(224),
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
        trans_after = data_transform(img_pil)
        return trans_after

    def classify_LJ(self,image):
        classes = classes_lj["CN"][0]
        model_path = classes_lj["CN"][1]
        num_classes = len(classes)  # 获取类别数量
        model = self.model
        model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.to(self.device)
        model.eval()
        img = self.date_transform(image)
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        img = img.to(self.device)
        outputs = model(img)
        out_put = torch.softmax(outputs, dim=1)
        maxk = max((1, 3))
        _, pred = outputs.topk(maxk, 1, True, True)
        label = classes[pred[0][0]]
        conf = out_put[0][pred[0][0]]
        torch.cuda.empty_cache()
        return label,conf

    def classify_NG(self,image,LJ_type):
        classes = classes_indices[LJ_type]
        num_classes = len(classes)
        model_path = classes_models[LJ_type]
        model = self.model
        model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.to(self.device)
        model.eval()
        img = self.date_transform(image)
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        img = img.to(self.device)
        outputs = model(img)
        out_put = torch.softmax(outputs, dim=1)
        maxk = max((1, 2))
        _, pred = outputs.topk(maxk, 1, True, True)
        label = classes[pred[0][0]]
        conf = out_put[0][pred[0][0]]
        torch.cuda.empty_cache()
        return label, conf






# models_dict = {}
# class_indices_dict = {}  # 将class_indices_dict声明为空字典
# for suffix, model_path in model_paths.items():
#     assert os.path.exists(model_path), f"文件'{model_path}'不存在"
#
#     class_indices_path = class_indices_paths[suffix]
#     assert os.path.exists(class_indices_path), f"文件'{class_indices_path}'不存在"
#     with open(class_indices_path, "r") as f:
#         class_indices = json.load(f)
#         class_indices_dict[suffix] = dict((val, key) for key, val in class_indices.items())
#         model = models.resnet34(pretrained=False)
#         num_classes = len(class_indices_dict[suffix])  # 获取类别数量
#         model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
#         model.load_state_dict(torch.load(model_path, map_location=device))
#         model.eval()
#         models_dict[suffix] = model.to(device)
#
# # 预测图片
# # OKjpg_path = os.path.join(rf"{add_file_name}/OKjpg")
# # OKptt_path = os.path.join(rf"{add_file_name}/OKptt")
# # check_folder_exit(rf"{add_file_name}/OKjpg/")
# # check_folder_exit(rf"{add_file_name}/OKptt/")
# # add_file_name = r'F:\1102\301\InspectResult\AI\6B01M004KV50T-V2(T)-FB\TOP\11'
# out_dir = os.path.join(add_file_name, "AIjpg")
# assert os.path.exists(out_dir), f"文件夹'{out_dir}'不存在"
#
# for suffix, model in models_dict.items():
#     if suffix not in ["_U", '_CN', '_X', '_J', '_PJ']:
#         continue
#     # class_dir = os.path.join(add_file_name, class_indices_dict[suffix])
#     class_dir = os.path.join(add_file_name, list(class_indices_dict[suffix].values())[0])
#
#     img_path_list = [os.path.join(out_dir, filename) for filename in os.listdir(out_dir) if
#                      filename.endswith(suffix + ".jpg")]
#     if len(img_path_list) == 0:
#         # print(f"文件夹'{out_dir}'下不存在后缀为'{suffix}.jpg'的文件")
#         continue
#     # class_indices_path = class_indices_paths[suffix]  # 根据后缀名获取对应的类别索引文件路径
#     # print(class_indices_path)
#     # model = models_dict[suffix]
#     # print(model)
#     class_indices_path = class_indices_paths[suffix]
#     model = models_dict[suffix]
#     print(class_indices_path)
#
#     txt_class_path = os.path.join(add_file_name, f"class.txt")
#     batch_size = 1
#     json_file = open(class_indices_path, "r")
#     class_indict = json.load(json_file)
#     with open(txt_class_path, "a", encoding="utf_8_sig") as f:
#         with torch.no_grad():
#             for ids in range(0, len(img_path_list) // batch_size):
#                 img_list = []
#                 for img_path in img_path_list[ids * batch_size: (ids + 1) * batch_size]:
#                     assert os.path.exists(img_path), f"file: '{img_path}' dose not exist."
#                     # img_basename = os.path.basename(img_path)
#                     # muban_path = os.path.join(muban_folder,img_basename)
#                     # muban_img = Image.open(muban_path)
#
#                     img = Image.open(img_path)
#                     img = data_transform(img)
#                     img_list.append(img)
#
#                 # batch img
#                 # 将img_list列表中的所有图像打包成一个batch
#                 batch_img = torch.stack(img_list, dim=0)
#                 # predict class
#                 output = model(batch_img.to(device)).cpu()
#                 predict = torch.softmax(output, dim=1)
#                 probs, classes = torch.max(predict, dim=1)
#                 ##
#                 # OKImage_path=os.path.join(rf"{add_file_name}/OKImage/")
#                 # check_folder_exit(rf"{add_file_name}/OKImage/")
#                 for idx, (pro, cla) in enumerate(zip(probs, classes)):
#                     # if pro.numpy()<=0.70:
#                     #     a='NG'
#                     # else:
#                     #    a=class_indict[str(cla.numpy())]
#
#                     print("image: {}  class: {}  prob: {:.3}".format(
#                         img_path_list[ids * batch_size + idx],
#                         class_indict[str(cla.numpy())],
#                         pro.numpy()), file=f)
#     # 分类图片到对应类别的文件夹下
#
#     AI = open(txt_class_path, "r", encoding='utf_8_sig')
#     for line in AI.readlines()[0:]:
#         if line != "":
#             new_path = line.split(" ")[4]
#             jpgpath = os.path.join(rf"{add_file_name}\{new_path}\2D")
#             pttpath = os.path.join(rf"{add_file_name}\{new_path}\3D")
#             jpg_path = line.split(" ")[1]
#             ptt_path = jpg_path.replace("AIjpg", "3D").replace('jpg', 'ptt')
#             check_folder_exit(jpgpath)
#             check_folder_exit(pttpath)
#             # Rename the images in jpgpath and pttpath
#             shutil.move(jpg_path, jpgpath)
#             shutil.move(ptt_path, pttpath)
#
#             file_name = re.search(r'[^\\/:*?"<>|\r\n]+$', jpg_path).group(0)
#
#             parts = file_name.rsplit("_", 1)
#
#             new_filename1 = parts[0] + '_' + new_path + '.' + 'jpg'
#
#             path1 = os.path.join(jpgpath, new_filename1)
#             os.rename(os.path.join(jpgpath, file_name), path1)
#             file_name = re.search(r'[^\\/:*?"<>|\r\n]+$', ptt_path).group(0)
#             new_filename2 = parts[0] + '_' + new_path + '.' + 'ptt'
#
#             path2 = os.path.join(pttpath, new_filename2)
#
#             os.rename(os.path.join(pttpath, file_name), path2)
#
#     print("IC,CN分类完成")
#     # ####### 加载模型和类别索引
#     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#     models_dict = {}
#     class_indices_dict = {}  # 将class_indices_dict声明为空字典
#     for suffix, model_path in model_paths.items():
#         assert os.path.exists(model_path), f"文件'{model_path}'不存在"
#
#         class_indices_path = class_indices_paths[suffix]
#         assert os.path.exists(class_indices_path), f"文件'{class_indices_path}'不存在"
#         with open(class_indices_path, "r") as f:
#             class_indices = json.load(f)
#             class_indices_dict[suffix] = dict(
#                 (val, key) for key, val in class_indices.items())
#         #
#         model = models.resnet34(pretrained=False)
#         num_classes = len(class_indices_dict[suffix])  # 获取类别数量
#         model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
#         model_path = model_paths[suffix]
#         model.load_state_dict(torch.load(model_path, map_location=device))
#         model.eval()
#         models_dict[suffix] = model.to(device)
#
#     # 预测图片
#
#     # AI1 = open(txt_class_path, "r")
#     AI.seek(0)
#     txt_path1 = os.path.join(add_file_name, f"yichang.txt")  ##异常类别分类
#     processed_paths = set()  # 维护已处理的 new_path 值的集合
#     for line in AI.readlines()[0:]:
#         if line != "":
#             new_path = line.split(" ")[4]
#             if new_path == "MIS" or new_path == "IC2" or new_path == "CN2":
#                 continue  # 如果 new_path 的值为 "MIS"或者"IC2"，则跳过后续操作
#             if new_path in processed_paths:
#                 continue  # 如果该路径已被处理，则跳过
#             processed_paths.add(new_path)  # 将该路径添加到已处理的集合中
#
#             out_dir = os.path.join(rf"{add_file_name}\{new_path}\2D")
#             assert os.path.exists(out_dir), f"文件夹'{out_dir}'不存在"
#
#             for suffix, model in models_dict.items():
#                 class_indices_path = class_indices_paths[suffix]
#                 model = models_dict[suffix]
#                 print(class_indices_path)
#                 # print(model)
#                 img_path_list = [os.path.join(out_dir, filename) for filename in
#                                  os.listdir(out_dir)
#                                  if filename.endswith(suffix + ".jpg")]
#                 if len(img_path_list) == 0:
#                     # print(f"文件夹'{out_dir}'下不存在后缀为'{suffix}.jpg'的文件")
#                     continue
#
#                 batch_size = 1  # 每次预测时将多少张图片打包成一个batch
#
#                 json_file = open(class_indices_path, "r")
#                 class_indict = json.load(json_file)
#
#                 with open(txt_path1, "a", encoding='utf_8_sig') as f:
#                     with torch.no_grad():
#                         for ids in range(0, len(img_path_list) // batch_size):
#                             img_list = []
#                             for img_path in img_path_list[
#                                             ids * batch_size: (ids + 1) * batch_size]:
#                                 assert os.path.exists(
#                                     img_path), f"file: '{img_path}' dose not exist."
#                                 img = Image.open(img_path)
#                                 img = data_transform(img)
#                                 img_list.append(img)
#
#                             # batch img
#                             # 将img_list列表中的所有图像打包成一个batch
#                             batch_img = torch.stack(img_list, dim=0)
#                             # predict class
#                             output = model(batch_img.to(device)).cpu()
#                             predict = torch.softmax(output, dim=1)
#                             probs, classes = torch.max(predict, dim=1)
#                             ##
#                             # OKImage_path=os.path.join(rf"{add_file_name}/OKImage/")
#                             # check_folder_exit(rf"{add_file_name}/OKImage/")
#                             for idx, (pro, cla) in enumerate(zip(probs, classes)):
#                                 # if pro.numpy()<=0.70:
#                                 #     a='NG'
#                                 # else:
#                                 #    a=class_indict[str(cla.numpy())]
#
#                                 print("image: {}  class: {}  prob: {:.3}".format(
#                                     img_path_list[ids * batch_size + idx],
#                                     class_indict[str(cla.numpy())],
#                                     pro.numpy()), file=f)